from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Callable

import casadi as ca


SymbolFactory = Callable[[str], ca.MX]


@dataclass(frozen=True)
class RobotModel:
    T: float
    N: int
    rob_diam: float
    v_max: float
    omega_max: float
    states: ca.MX
    controls: ca.MX
    rhs: ca.MX
    f: ca.Function
    n_states: int
    n_controls: int


def build_robot_model(
    symbol_factory: SymbolFactory = ca.SX.sym,
    *,
    T: float = 0.2,
    N: int = 100,
    rob_diam: float = 0.3,
    v_max: float = 0.6,
    omega_max: float = math.pi / 4.0,
) -> RobotModel:
    """
    Return a configurable robot model setup used across the MPC examples.
    """
    x = symbol_factory("x")
    y = symbol_factory("y")
    theta = symbol_factory("theta")
    states = ca.vertcat(x, y, theta)

    v = symbol_factory("v")
    omega = symbol_factory("omega")
    controls = ca.vertcat(v, omega)

    rhs = ca.vertcat(v * ca.cos(theta), v * ca.sin(theta), omega)
    f = ca.Function(
        "f",
        [states, controls],
        [rhs],
        ["input_state", "control_input"],
        ["rhs"],
    )

    return RobotModel(
        T=T,
        N=N,
        rob_diam=rob_diam,
        v_max=v_max,
        omega_max=omega_max,
        states=states,
        controls=controls,
        rhs=rhs,
        f=f,
        n_states=int(states.numel()),
        n_controls=int(controls.numel()),
    )
